from BitVector import BitVector
from random import Random
import math

#
# ryan.a.cox@gmail.com / www.asciiarmor.com
#
# copyright (c) 2008, ryan cox
# all rights reserved 
# BSD license: http://www.opensource.org/licenses/bsd-license.php
#

class BloomFilter(object):
	def __init__(self, m=None, k=None, n=None, p=None, bits=None ):
		if k and (k > 4 or k < 1):
			raise Exception('Must specify value of k between 1 and 4')

		# create filter from existing bitvector / number of bits in array / hash count
		if bits and (m and k) and (not n) and (not p):	
			self.bits = bits
			self.m = m
			self.k = k
		# create filter based on number of bits in array / hash count
		elif (m and k) and (not p) and (not n) and (not bits):
			self.bits = BitVector( size=m )
			self.m = m
			self.k = k
		# create filter based on max number elements / acceptable error rate 
		elif (not m and not k) and (p and n) and (not bits):
			self.m = int(math.ceil((n * math.log(p)) / math.log(1.0 / (math.pow(2.0, math.log(2.0))))))
			self.k = int(math.ceil(math.log(2.0) * self.m / n))
			self.n = n
			self.p = p
			self.bits = BitVector( size=self.m )
		else:
			raise Exception('Illegal construction. Must specify specify m / k or m / k / bits or n /p ')
		self.rand = Random()
	
	def __contains__(self, key):
		for i in self._indexes(key): 
			if not self.bits[i]:
				return False	
		return True 

	def add(self, key):
		dupe = True 
		bits = []
		for i in self._indexes(key): 
			if dupe and not self.bits[i]:
				dupe = False
			self.bits[i] = 1
			bits.append(i)
		return dupe

	def __eq__(self, f):
		return self.k == f.k and self.m == f.m and self.bits == f.bits

	def __ne__(self, f):
		return not self.__eq__(f)

	def __and__(self, filter):
		if (self.k != filter.k) or (self.m != filter.m): 
			raise Exception('Must use bloom filters created with equal k / m paramters for bitwise AND')
		return BloomFilter(m=self.m,k=self.k,bits=(self.bits & filter.bits))

	def __or__(self, filter):
		if (self.k != filter.k) or (self.m != filter.m): 
			raise Exception('Must use bloom filters created with equal k / m paramters for bitwise OR')
		return BloomFilter(m=self.m,k=self.k,bits=(self.bits | filter.bits))

	def _indexes(self,key):
		self.rand.seed(hash(key))
		ret = []
		for i in range(self.k):
			ret.append(self.rand.randint(0,self.m-1))
		return ret

if __name__ == '__main__':

	# test creation from element count / acceptable probability
	t = BloomFilter(n=3000, p=0.01)
	assert t.m == 28756
	assert t.k == 7

	# test equals / not equals
	a = BloomFilter(m=100, k=4)
	b = BloomFilter(m=100, k=4)
	assert a == b
	a.add('1')
	a.add('2')
	a.add('3')
	b.add('1')
	b.add('2')
	b.add('3')
	assert a == b
	b.add('4')
	assert a != b

	# create test filters
	e = BloomFilter(m=100, k=4)
	f = BloomFilter(m=100, k=4)
	
	e.add('one')
	e.add('two')
	e.add('three')
	e.add('four')
	e.add('five')

	f.add('three')
	f.add('four')
	f.add('five')
	f.add('six')
	f.add('seven')
	f.add('eight')
	f.add('nine')
	f.add("ten")

	# test check for dupe on add
	assert not f.add('eleven') 
	assert f.add('eleven') 
	
	# test membership operations
	assert 'ten' in f 
	assert 'one' in e 
	assert 'ten' not in e 
	assert 'one' not in f 

	# test set based operations
	union = f | e
	intersection = f & e

	assert 'ten' in union
	assert 'one' in union 
	assert 'three' in intersection
	assert 'ten' not in intersection
	assert 'one' not in intersection

	i = 0
	for word in open('american-english'):
		i += 1
	print "loading %d words..." % i
	
	# test creation from element count / acceptable probability
	dict = BloomFilter(n=i, p=0.01)

	collissions = 0
	for word in open('american-english'):
		if dict.add(word.strip()):
			collissions += 1
	print "loaded %d words into array %d bits in length using %d hashes" % ( i,dict.m,dict.k)
	print "%d collissions ( %.2f%% )" % (collissions,(collissions/float(i))*100)

	# simple spell check tests
	assert 'zone' in dict
	assert 'untimeliness' in dict
	assert 'dalmatians' in dict
	assert 'dalmations' not in dict
	assert 'xxxyz' not in dict

